#include <cv.h>
#include <highgui.h>
#include <stdio.h>
#include <algorithm>
#include <iostream>
#include <map>
#include "contours.h"

#define BIN_THRES 75

//filtering of blobs
#define MIN_AREA 1000

//classification threshold
#define COMPACTNESS_MARGIN 0.08

cv::Mat origImg;
cv::Mat binImg;

// points and contour are in image coords, but pointMat, and therefore the moments,
// are relative to offset
struct Blob {
	cv::Point offset;
	std::vector<cv::Point> points;
	cv::Mat pointMat; //for each point in points, this has a 255
	std::vector<cv::Point> contour;
	double area;
	double circum;
	cv::Moments m;
};

//get min/max distance on contour relative to center
std::vector<double> getMinMax(cv::Point c, std::vector<cv::Point> contour) {
	assert(contour.size() >= 1);
	double minDist = cv::norm(c-contour[0]);
	double maxDist = minDist;
	uint minInd = 0;
	uint maxInd = 0;
	
	for (uint i = 1; i < contour.size(); i++) {
		double diff = cv::norm(c-contour[i]);
		if (diff < minDist) {
			minDist = diff;
			minInd = i;
		} else if (diff > maxDist) {
			maxDist = diff;
			maxInd = i;
		}
	}
	
	std::vector<double> minMax(2);
	minMax[0] = minDist;
	minMax[1] = maxDist;
	return minMax;
}

int indexOf(std::vector<int>& v, int e) {
	int index = -1;
	for (uint i = 0; i < v.size(); i++) {
		if (v[i] == e) {
			index = i;
			break;
		}
	}
	return index;
}

int resolveMarker(std::vector<int>& m, int i) {
	if (i == m[i]) {
		return i;
	} else {
		return m[i] = resolveMarker(m, m[i]); //should reduce recursion depth for subsequent calls
	}
}

cv::Mat getBinaryImage(cv::Mat& img, double thres) {
	//the binarization itself doesn't actually convert the image to greyscale, so it has to be done by hand
	cv::Mat tmpGreyImg;
	cv::cvtColor(img, tmpGreyImg, CV_BGR2GRAY);

	//I assume potential ellipses should be white, hence the _INV
	cv::threshold(tmpGreyImg, binImg, thres, 255, cv::THRESH_BINARY_INV);
	return binImg;
}

std::vector<cv::Point> getContour(cv::Mat pointMat) {
	//TODO: the algorithm in the pdf instead of the opencv call?
	std::vector<cv::Point> contour = v7::findContour(pointMat);
	/*
	std::vector<std::vector<cv::Point> > contours;
	cv::findContours(pointMat.clone(), contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_NONE);
	
	assert(contours.size() == 1); //something went wrong if this doesnt hold true
	return contours[0];
	*/
	return contour;
}

std::vector<Blob> getBlobs(cv::Mat& image) {
	cv::Mat_<int> img(image); //grayscale image assumed, conversion need for easy access
	//first will be projected to second
	std::vector<int> eqMap;
	eqMap.push_back(0); //so [0] == 0
	cv::Mat_<int> markerMat(img.rows, img.cols);
	
	int marker = 1;
	for (int y = 0; y < markerMat.rows; y++) {
		for (int x = 0; x < markerMat.cols; x++) {
			int a = img(y, x);
			int l = 0;
			int o = 0;
			int mo = 0;
			int ml = 0;
			
			if (a < 255) {
				markerMat(y, x) = 0;
				continue;			
			}
			
			if (x > 0) {
				l = img(y, x - 1);
				ml = resolveMarker(eqMap, markerMat(y, x - 1));
			}
			if (y > 0) {
				o = img(y - 1,  x);
				mo = resolveMarker(eqMap, markerMat(y - 1, x));
			}
			
			if (a == o && a != l) {
				markerMat(y, x) = mo;
			} else if (a != o && a == l) {
				markerMat(y, x) = ml;
			} else if (a == o && a == l) { 
				markerMat(y, x) = ml; //arbitrary
				eqMap[mo] = ml; //if mo==ml, this does nothing
			} else {
				markerMat(y, x) = marker;
				eqMap.push_back(marker);
				marker++; 
			}
		}
	}
	
	//get list of unique markers
	std::vector<int> uniqueMarkers;
	
	for (uint i = 0; i < eqMap.size(); i++) {
		int index = indexOf(uniqueMarkers, eqMap[i]);
		if (index < 0) uniqueMarkers.push_back(eqMap[i]);
	}

	//divide marker areas into point lists, this also serves to isolate each blob from others (else the moments would be a bit wrong)
	std::vector<std::vector<cv::Point> > pointList(uniqueMarkers.size());
	for (int y = 0; y < markerMat.rows; y++) {
		for (int x = 0; x < markerMat.cols; x++) {
			if (markerMat(y,x) == 0) continue;
			int markerIndex = -1;
			int resolvedMarker = resolveMarker(eqMap, markerMat(y, x));
			for (uint i = 0; i < uniqueMarkers.size(); i++) {
				if (resolvedMarker == uniqueMarkers[i]) {
					markerIndex = i;
					break;
				}
			}
			pointList[markerIndex].push_back(cv::Point(x, y));
		}	
	}
	
	//get blob list
	std::vector<Blob> blobList;
	for (uint i = 0; i < pointList.size(); i++) {
		if (pointList[i].size() < MIN_AREA) continue;
		Blob b;
		b.points = pointList[i];
		b.area = pointList[i].size();
		
		cv::Rect bb = cv::boundingRect(b.points);
		b.pointMat = cv::Mat(bb.height, bb.width, CV_8UC1);
		for (int y = 0; y < b.pointMat.rows; y++) {
			for (int x = 0; x < b.pointMat.cols; x++) {
				b.pointMat.at<uchar>(y, x) = 0;
			}
		}
		
		for (uint j = 0; j < b.points.size(); j++) {
			b.pointMat.at<uchar>(b.points[j].y - bb.y, b.points[j].x - bb.x) = 255;
		}
		
		b.offset = cv::Point(bb.x, bb.y);
		b.contour = getContour(b.pointMat);
		for (uint j = 0; j < b.contour.size(); j++) {
			b.contour[j] += cv::Point(bb.x, bb.y);
		}
		
		b.circum = b.contour.size();
		b.m = cv::moments(b.pointMat, true);
		blobList.push_back(b);
	}
	return blobList;
}

void findEllipses(cv::Mat& original) {
	cv::Mat binImg = getBinaryImage(original, BIN_THRES);

	std::vector<Blob> blobs = getBlobs(binImg);
	for (uint i = 0; i < blobs.size(); i++) {
		Blob blob = blobs[i];

		cv::Moments m = blob.m;
		double x = m.m10/m.m00;
		double y = m.m01/m.m00;
		
		std::vector<double> minMax = getMinMax(cv::Point(x, y) + blob.offset, blob.contour); //calc ellipse from minmax
		double l1 = minMax[0];
		double w1 = minMax[1];
		
		double calcArea1 = M_PI * l1 * w1;
		double calcCirc1 = M_PI * (3 * (l1+w1) - sqrt((3*l1+w1)*(l1+3*w1))); //ramanujan approx for circumference
		double calcCompact1 = calcCirc1 * calcCirc1 / (4 * M_PI * calcArea1);
		
		cv::RotatedRect rr = cv::minAreaRect(blob.points); //get a minimal, rotatable rect that encloses the ellipse
		double l2 = std::min(rr.size.width, rr.size.height) / 2.0;
		double w2 = std::max(rr.size.width, rr.size.height) / 2.0;
		
		double calcArea2 = M_PI * l2 * w2;
		double calcCirc2 = M_PI * (3 * (l2+w2) - sqrt((3*l2+w2)*(l2+3*w2)));
		double calcCompact2 = calcCirc2 * calcCirc2 / (4 * M_PI * calcArea2);
		
		//technically, the comparison of both l and w should be enough
		double compactDiff = fabs(calcCompact2 - calcCompact1); 
		if (compactDiff < COMPACTNESS_MARGIN) {
			std::vector<std::vector<cv::Point> > contours;
			contours.push_back(blob.contour);
			cv::drawContours(origImg, contours, 0, cv::Scalar(0, 0, 255, 0));
			printf("found ellipse at (%.1f, %.1f) with confidence %.3f\n", blob.offset.x + x, blob.offset.y + y, 1 - compactDiff);
		}
	}
}

int main(int argc, char **argv) {
	//origImg = cv::imread(argv[1], CV_LOAD_IMAGE_COLOR);
	// given an IplImage* iimg, build a Mat via: "Mat(iimg)"
	for (int i = 1; i < argc; i++) {
		origImg = cv::imread(argv[i], CV_LOAD_IMAGE_COLOR);
		
		findEllipses(origImg);

		//cv::namedWindow("Binary Image", CV_WINDOW_AUTOSIZE);
		//cv::imshow("Binary Image", binImg);
		cv::namedWindow("Found Ellipses", CV_WINDOW_AUTOSIZE);
    	cv::imshow("Found Ellipses", origImg);
	
   		cv::waitKey(0);
	}
	return 0;
}
